package org.study.dsm.mybatisplus.config;

import com.alibaba.druid.pool.DruidDataSource;
import com.alibaba.druid.stat.DataSourceMonitorable;
import com.alibaba.druid.stat.DruidDataSourceStatManager;
import com.google.common.collect.Maps;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.study.dsm.mybatisplus.model.DataSourceInfo;

import java.sql.*;
import java.util.Map;
import java.util.Set;

/**
 * 动态数据源
 * 调用AddDefineDataSource组件的addDefineDynamicDataSource（）方法，获取原来targetdatasources的map，并将新的数据源信息添加到map中，并替换targetdatasources中的map
 * 切换数据源时可以使用@DataSource(value = "数据源名称")，或者DynamicDataSourceContextHolder.setContextKey("数据源名称")
 *
 * @author doveylovey
 * @version v1.0.0
 * @email 1135782208@qq.com
 * @date 2024年04月27日
 */
@Slf4j
public class DynamicDataSource extends AbstractRoutingDataSource {
    @Getter
    private Map<Object, Object> dynamicTargetDataSources = Maps.newConcurrentMap();

    /**
     * 检测并创建动态数据源：数据源不存在则创建，存在则创建
     *
     * @param dataSourceInfo 数据源信息
     */
    public boolean createDataSourceWithCheck(DataSourceInfo dataSourceInfo) {
        String dsKey = dataSourceInfo.getDsKey();
        if (this.dynamicTargetDataSources.containsKey(dsKey)) {
            log.info("数据源[{}]之前已经创建，准备测试数据源是否正常", dsKey);
            DruidDataSource druidDataSource = (DruidDataSource) this.dynamicTargetDataSources.get(dsKey);
            log.info("[{}]数据源的概况：{}个闲置连接, {}个活动连接", dsKey, druidDataSource.getPoolingCount(), druidDataSource.getActiveCount());
            boolean rightFlag = true;
            try (Connection connection = druidDataSource.getConnection()) {
                log.info("数据源[{}]正常", dsKey);
            } catch (Exception e) {
                rightFlag = false;
                log.error(e.getMessage());
                log.info("缓存数据源[{}]已失效，准备删除", dsKey);
                if (this.deleteDataSources(dsKey)) {
                    log.info("缓存数据源删除成功");
                } else {
                    log.error("缓存数据源删除失败");
                }
            }
            if (rightFlag) {
                return true;
            } else {
                return this.createDataSource(dataSourceInfo);
            }
        } else {
            return this.createDataSource(dataSourceInfo);
        }
    }

    /**
     * 删除数据源
     *
     * @param dsKey 数据源key
     */
    public boolean deleteDataSources(String dsKey) {
        if (this.dynamicTargetDataSources.containsKey(dsKey)) {
            Set<DataSourceMonitorable> dataSourceMonitorableSet = DruidDataSourceStatManager.getDruidDataSourceInstances();
            for (DataSourceMonitorable monitorable : dataSourceMonitorableSet) {
                if (dsKey.equals(monitorable.getName())) {
                    this.dynamicTargetDataSources.remove(dsKey);
                    DruidDataSourceStatManager.removeDataSource(monitorable);
                    // 将map赋值给父类的TargetDataSources
                    super.setTargetDataSources(this.dynamicTargetDataSources);
                    // 将TargetDataSources中的连接信息放入resolvedDataSources管理
                    super.afterPropertiesSet();
                    return true;
                }
            }
        }
        return false;
    }

    /**
     * 查看数据源是否存在
     *
     * @param dsKey 数据源key
     */
    public boolean existsDataSource(String dsKey) {
        return this.dynamicTargetDataSources.containsKey(dsKey);
    }

    /**
     * 检测数据源是否正确
     *
     * @param dsKey 数据源key
     */
    public boolean checkDataSource(String dsKey) {
        if (this.existsDataSource(dsKey)) {
            DruidDataSource druidDataSource = (DruidDataSource) this.dynamicTargetDataSources.get(dsKey);
            log.info("[{}]数据源的概况：{}个闲置连接, {}个活动连接", dsKey, druidDataSource.getPoolingCount(), druidDataSource.getActiveCount());
            try (Connection connection = druidDataSource.getConnection();
                 Statement statement = connection.createStatement()) {
                statement.execute("SELECT 1");
                return true;
            } catch (SQLException e) {
                return false;
            }
        }
        return false;
    }

    /**
     * 创建数据源
     *
     * @param dataSourceInfo 数据源信息
     */
    private boolean createDataSource(DataSourceInfo dataSourceInfo) {
        String dsKey = dataSourceInfo.getDsKey();
        log.info("准备创建数据源：[{}]", dsKey);
        String username = dataSourceInfo.getUsername();
        String password = dataSourceInfo.getPassword();
        String url = dataSourceInfo.getUrl();
        String driveClass = "com.mysql.cj.jdbc.Driver";// dataSourceInfo.getDrive();
        try {
            Class.forName(driveClass);
            DriverManager.getConnection(url, username, password);
        } catch (Exception ex) {
            log.info("测试数据源是否有效异常", ex);
            return false;
        }
        try {
            Class.forName(driveClass);
            // 相当于连接数据库
            DriverManager.getConnection(url, username, password);
            // 创建数据源
            DruidDataSource druidDataSource = new DruidDataSource();
            druidDataSource.setName(dsKey);
            druidDataSource.setDriverClassName(driveClass);
            druidDataSource.setUrl(url);
            druidDataSource.setUsername(username);
            druidDataSource.setPassword(password);
            druidDataSource.init();
            this.dynamicTargetDataSources.put(dsKey, druidDataSource);
            // 将map赋值给父类的TargetDataSources
            this.setTargetDataSources(this.dynamicTargetDataSources);
            log.info("数据源[{}]初始化成功", dsKey);
            return true;
        } catch (Exception ex) {
            log.error("创建数据源异常", ex);
        }
        return false;
    }

    /**
     * 设置当前可以使用的数据源，key是数据源的唯一表示，value是具体数据源
     * 该方法是在动态数据源被初始化后调用的(在当前类中)，而不是在数据源切换时被调用，因为数据源切换有两种情况：
     * 1、数据源没有被初始化，那么会先初始化数据源，然后修改DynamicDataSourceHolder
     * 2、数据源已经被初始化，那么会直接修改DynamicDataSourceHolder
     */
    @Override
    public void setTargetDataSources(Map<Object, Object> targetDataSources) {
        // 通过super.setTargetDataSources()将当前已经存在的数据源Map集传递给AbstractRoutingDataSource的targetDataSources
        super.setTargetDataSources(targetDataSources);
        // 利用afterPropertiesSet()将targetDataSources赋值给resolvedDataSources
        super.afterPropertiesSet();
        this.dynamicTargetDataSources = targetDataSources;
    }

    /**
     * 该方法是实现动态数据源切换的核心方法，是在AbstractRoutingDataSource中被调用的，决定当前线程使用哪个数据源
     * 创建的数据源都是存储在Map中的，determineCurrentLookupKey()的作用就是拿到目标数据源的key，从而获取到目标数据源
     */
    @Override
    protected Object determineCurrentLookupKey() {
        String dsKey = DynamicDataSourceHolder.getDataSource();
        if (StringUtils.isNotBlank(dsKey)) {
            if (!this.dynamicTargetDataSources.containsKey(dsKey)) {
                throw new RuntimeException("数据源不存在");
            }
            log.info("当前数据源：[{}]", dsKey);
        } else {
            log.info("当前数据源：[默认数据源]");
        }
        return dsKey;
    }
}
